from time import sleep
from zlib import compress
from threading import Thread
from socket import (SO_RCVLOWAT, socket, AF_INET, SOCK_DGRAM, SOL_SOCKET, SO_BROADCAST)
from tkinter import Tk, BooleanVar, Button, Label, Tk
from PIL.ImageGrab import grab

PORT = 22222

root = Tk()
root.title("屏幕广播程序")
# 初始化窗口尺寸
root.geometry('320x60+500+200')
# 两个方向都不允许缩放
root.resizable(False, False)

# 缓冲区大小
BUFFER_SIZE = 1400
# 控制是否正在发送的变量
sending = BooleanVar(root, value=False)

def send_image():
    sock = socket(AF_INET, SOCK_DGRAM)
    # 广播
    sock.setsockopt(SOL_SOCKET, SO_BROADCAST, True)
    IP = "255.255.255.255"

    while sending.get():
        # 屏幕截图
        im = grab()
        w, h = im.size
        # 把图片转换为字节，压缩后传输，减少带宽占用
        im_bytes = compress(im.tobytes())
        # 通知局域网内用户开始发送
        sock.sendto(b'start', (IP, PORT))

        # 分块发送，以免字节太长无法发送
        print("Start sending")
        for i in range(len(im_bytes) // BUFFER_SIZE + 1):
            start = i * BUFFER_SIZE
            end =  start + BUFFER_SIZE
            sock.sendto(im_bytes[start:end], (IP, PORT))
            sleep(0.0002)
        # 发送完毕
        print("Send stop")
        sock.sendto(b'_over' + str((w, h)).encode(), (IP, PORT))

        sleep(0.5) # 0.5 秒后继续截图发送
    sock.sendto(b'close', (IP, PORT))
    sock.close()

lbCopyRight = Label(root, text='网络编程课程设计实验二，基于UDP的屏幕发送程序', fg='red', cursor='plus')
lbCopyRight.place(x=5, y=5, width=310, height=20)

def btnStartClick():
    sending.set(True)
    Thread(target=send_image).start()
    btnStart['state'] = 'disable'
    btnStop['state'] = 'normal'
btnStart = Button(root, text="开始广播", command=btnStartClick)
btnStart.place(x=30, y=30, width=125, height=20)

def btnStopClick():
    sending.set(False)
    btnStart['state'] = 'normal'
    btnStop['state'] = 'disable'
btnStop = Button(root, text='停止广播', command=btnStopClick)
btnStop['state'] = 'disable'
btnStop.place(x=165, y=30, width=125, height=20)

root.mainloop()
